import base64
import os
from reportlab.lib import colors
from reportlab.lib.pagesizes import A4
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Image, Table, TableStyle
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
from reportlab.pdfbase import pdfmetrics
from reportlab.pdfbase.ttfonts import TTFont
from io import BytesIO
import ujson as json
import datetime
from data_utils import format_calc_id

# 使用系统默认中文字体（示例：Windows 上的 SimSun）
SYSTEM_FONT_PATH = "fonts/SourceHanSansSC-VF.ttf"
pdfmetrics.registerFont(TTFont('source-hans-sc', SYSTEM_FONT_PATH))

def create_calc_pdf(calc_name, inputs, result, timestamp=None, calc_id=None, calc_version=None, task_id=None, username=None):
    """生成计算结果PDF，优化排版和样式"""
    buffer = BytesIO()

    # 创建带有页眉的PDF文档
    class PDFWithHeader(SimpleDocTemplate):
        def __init__(self, *args, **kwargs):
            self.header_info = kwargs.pop('header_info', {})
            SimpleDocTemplate.__init__(self, *args, **kwargs)

        def afterPage(self):
            canvas = self.canv
            # 设置页脚字体和大小
            canvas.setFont('source-hans-sc', 8)

            # 在页面底部绘制页脚信息（单行，左对齐）
            footer_elements = []
            if self.header_info.get('calc_id'):
                footer_elements.append(f"计算器ID: {format_calc_id(self.header_info['calc_id'])}")
            if self.header_info.get('calc_version'):
                footer_elements.append(f"版本号: {self.header_info['calc_version']}")
            if self.header_info.get('task_id'):
                footer_elements.append(f"任务编号: {self.header_info['task_id']}")
            if self.header_info.get('username'):
                footer_elements.append(f"计算用户: {self.header_info['username']}")

            # 将所有元素用分隔符连接成一行
            footer_text = " | ".join(footer_elements)
            canvas.drawString(self.leftMargin, 20, footer_text)  # 20是距离底部的高度

    # 创建文档实例，传入页眉信息
    header_info = {
        'calc_id': calc_id,
        'calc_version': calc_version,
        'task_id': task_id,
        'username': username
    }

    doc = PDFWithHeader(
        buffer,
        pagesize=A4,
        rightMargin=50,
        leftMargin=50,
        topMargin=100,  # 增加顶部边距，为页眉留出空间
        bottomMargin=50,
        header_info=header_info
    )

    story = []
    styles = getSampleStyleSheet()

    # 优化标题样式
    styles.add(ParagraphStyle(
        name='ChineseTitle',
        fontName='source-hans-sc',
        fontSize=24,
        leading=32,
        spaceAfter=30,
        alignment=1,
        textColor=colors.HexColor('#2C3E50')  # 深蓝色标题
    ))

    # 优化子标题样式
    styles.add(ParagraphStyle(
        name='ChineseHeading1',
        fontName='source-hans-sc',
        fontSize=18,
        leading=25,
        spaceAfter=15,
        textColor=colors.HexColor('#34495E'),  # 稍浅的蓝色
        borderPadding=10,
        borderWidth=0,
        borderColor=colors.HexColor('#ECF0F1'),  # 淡灰色背景
        backColor=colors.HexColor('#ECF0F1')
    ))

    # 优化正文样式
    styles.add(ParagraphStyle(
        name='ChineseBody',
        fontName='source-hans-sc',
        fontSize=12,
        leading=20,
        spaceAfter=12,
        textColor=colors.HexColor('#2C3E50')
    ))

    # 添加标题和时间戳
    title = f"<b>{calc_name}</b>"
    story.append(Paragraph(title, styles['ChineseTitle']))
    if timestamp:
        story.append(Paragraph(f"计算时间：{timestamp}", styles['ChineseBody']))
    story.append(Spacer(1, 20))

    # 添加输入参数表格
    story.append(Paragraph("输入参数", styles['ChineseHeading1']))
    story.append(Spacer(1, 10))

    # 优化表格样式 - 直接使用转换后的title作为显示
    input_data = [["参数名", "参数值"]] + [[k, str(v)] for k, v in inputs.items()]
    input_table = Table(input_data, colWidths=[200, 250])
    input_table.setStyle(TableStyle([
        ('FONT', (0, 0), (-1, -1), 'source-hans-sc'),
        ('BACKGROUND', (0, 0), (-1, 0), colors.HexColor('#34495E')),  # 表头背景色
        ('TEXTCOLOR', (0, 0), (-1, 0), colors.white),  # 表头文字颜色
        ('ALIGN', (0, 0), (-1, -1), 'CENTER'),
        ('GRID', (0, 0), (-1, -1), 1, colors.HexColor('#BDC3C7')),  # 淡灰色网格线
        ('PADDING', (0, 0), (-1, -1), 8),  # 单元格内边距
        ('BACKGROUND', (0, 1), (-1, -1), colors.white),  # 数据行背景色
        ('ROWBACKGROUNDS', (0, 1), (-1, -1), [colors.white, colors.HexColor('#F9F9F9')])  # 斑马纹
    ]))
    story.append(input_table)
    story.append(Spacer(1, 20))

    # 添加计算结果
    story.append(Paragraph("计算结果", styles['ChineseHeading1']))
    story.append(Spacer(1, 10))

    if isinstance(result, dict):
        if result['type'] == 'single':
            story.append(Paragraph(str(result['value']), styles['ChineseBody']))
        elif result['type'] == 'multiple':
            for i, value in enumerate(result['values'], 1):
                story.append(Paragraph(f"● {value}", styles['ChineseBody']))

        if 'plot' in result:
            story.append(Spacer(1, 15))
            img_data = base64.b64decode(result['plot'])
            img = Image(BytesIO(img_data))
            # 计算宽高比并保持原比例
            aspect_ratio = img.imageWidth / img.imageHeight
            draw_width = 450
            draw_height = draw_width / aspect_ratio
            img.drawWidth = draw_width
            img.drawHeight = draw_height
            story.append(img)
    else:
        story.append(Paragraph(str(result), styles['ChineseBody']))

    doc.build(story)
    buffer.seek(0)
    return buffer

def create_calculator_code_pdf(calculator, timestamp=None):
    """生成计算器代码和配置的PDF文档

    Args:
        calculator: 计算器对象，包含名称、代码、配置等信息
        timestamp: 可选，导出时间戳
    """
    buffer = BytesIO()

    # 创建PDF文档
    doc = SimpleDocTemplate(
        buffer,
        pagesize=A4,
        rightMargin=50,
        leftMargin=50,
        topMargin=50,
        bottomMargin=50,
    )

    story = []
    styles = getSampleStyleSheet()

    # 优化标题和正文样式
    styles.add(ParagraphStyle(
        name='ChineseTitle',
        fontName='source-hans-sc',
        fontSize=24,
        leading=32,
        spaceAfter=30,
        alignment=1,
        textColor=colors.HexColor('#2C3E50')
    ))

    styles.add(ParagraphStyle(
        name='ChineseHeading',
        fontName='source-hans-sc',
        fontSize=16,
        leading=25,
        spaceAfter=15,
        textColor=colors.HexColor('#34495E')
    ))

    styles.add(ParagraphStyle(
        name='ChineseBody',
        fontName='source-hans-sc',
        fontSize=12,
        leading=20,
        spaceAfter=12,
        textColor=colors.HexColor('#2C3E50')
    ))

    # 添加标题
    story.append(Paragraph(f"{calculator.name}", styles['ChineseTitle']))
    story.append(Spacer(1, 20))

    # 添加基本信息
    basic_info = [
        (f"计算器ID: {format_calc_id(calculator.id)}"),
        (f"计算器版本号: {calculator.version}"),
        (f"最后保存时间: {calculator.updated_at.strftime('%Y-%m-%d %H:%M:%S')}"),
        (f"导出时间: {timestamp or datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    ]

    for info in basic_info:
        story.append(Paragraph(info, styles['ChineseBody']))
    story.append(Spacer(1, 20))

    # 添加计算器代码
    story.append(Paragraph("计算器代码", styles['ChineseHeading']))
    story.append(Spacer(1, 10))

    # 使用等宽字体显示代码
    code_style = ParagraphStyle(
        'Code',
        parent=styles['ChineseBody'],
        fontSize=10,
        leading=12,
        spaceAfter=20
    )

    # 将代码按行分割并添加行号
    code_lines = calculator.code.split('\n')
    formatted_code = []
    for i, line in enumerate(code_lines, 1):
        # 转义特殊字符
        line = line.replace('&', '&amp;').replace('<', '&lt;').replace('>', '&gt;')
        # 添加行号，并在每行末尾添加换行符
        formatted_code.append(f"{i:3d} | {line}<br/>")

    story.append(Paragraph('<pre>' + ''.join(formatted_code) + '</pre>', code_style))
    story.append(Spacer(1, 20))

    # 添加输入参数配置
    story.append(Paragraph("输入参数配置", styles['ChineseHeading']))
    story.append(Spacer(1, 10))

    # 创建输入参数表格
    if calculator.inputs:
        # 获取所有可能的列
        all_columns = set()
        for param in calculator.inputs:
            all_columns.update(param.keys())

        # 创建表头
        columns = ['name', 'title', 'type', 'target'] + [col for col in all_columns if col not in ['name', 'title', 'type', 'target']]
        table_data = [columns]  # 表头行

        # 添加数据行
        for param in calculator.inputs:
            row = [str(param.get(col, '')) for col in columns]
            table_data.append(row)

        # 创建表格并设置样式，减小列宽和字体
        col_widths = [60] * len(columns)  # 减小列宽
        table = Table(table_data, colWidths=col_widths)
        table.setStyle(TableStyle([
            ('FONT', (0, 0), (-1, -1), 'source-hans-sc'),
            ('FONTSIZE', (0, 0), (-1, -1), 8),  # 减小字体大小
            ('BACKGROUND', (0, 0), (-1, 0), colors.HexColor('#34495E')),
            ('TEXTCOLOR', (0, 0), (-1, 0), colors.white),
            ('ALIGN', (0, 0), (-1, -1), 'CENTER'),
            ('GRID', (0, 0), (-1, -1), 1, colors.HexColor('#BDC3C7')),
            ('PADDING', (0, 0), (-1, -1), 4),  # 减小内边距
            ('BACKGROUND', (0, 1), (-1, -1), colors.white),
            ('ROWBACKGROUNDS', (0, 1), (-1, -1), [colors.white, colors.HexColor('#F9F9F9')])
        ]))
        story.append(table)

    # 构建PDF
    doc.build(story)
    buffer.seek(0)
    return buffer

def generate_pdf_download_button(st, task_id, button_text="下载 PDF", position=""):
    """生成PDF下载按钮，从数据库读取任务信息

    Args:
        st: Streamlit对象
        task_id: 任务ID
        button_text: 按钮文字
        position: 按钮位置标识符(用于生成唯一key)
    """
    from data_utils import load_task_by_id

    # 从数据库加载任务信息
    task_info = load_task_by_id(task_id)
    if not task_info:
        st.error(f"未找到任务记录: {task_id}")
        return

    # 生成PDF
    pdf_buffer = create_calc_pdf(
        calc_name=task_info['calculator_name'],
        inputs=task_info['inputs'],
        result=task_info['output'],
        timestamp=task_info['timestamp'],
        calc_id=task_info['calculator_id'],
        calc_version=task_info['calculator_version'],
        task_id=task_info['task_id'],
        username=task_info['username']
    )

    # 生成文件名
    filename = f"{task_info['calculator_name']}_CT{task_id}_{task_info['timestamp'].replace(' ', '_').replace(':', '-')}.pdf"

    # 创建下载按钮
    st.download_button(
        label=button_text,
        data=pdf_buffer,
        file_name=filename,
        mime="application/pdf",
        key=f"pdf_download_{task_id}{position}"
    )
